#include <GL/glew.h>
#include <GL/freeglut.h>

#include "cuda.h"
#include "cuda_runtime.h"
#include "cuda_gl_interop.h"
#include "cutil_math.h"
#include <cutil_inline.h>   


#include <cstdlib>
#include <cstdio>
#include <cassert>


#include "fluid.h"
#include "shader.h"

#define JITERATION 2
#define TRACE_COLOR (make_float4(h_color_r , h_color_g , h_color_b , 1.f)*injection_rate)


//DEBUG
bool vel_diffusion_enable =true;
bool vel_advect_enable = true;
bool vel_projection_enable = true;
bool density_diffusion_enable = true;
bool density_advect_enable = true;
bool haptic_effect_enable = true;


///////////////////////////
#define CUDA_SAFE_CALL(expr)						\
  {									\
    cudaError_t err = expr;						\
    if (err != cudaSuccess)						\
      {									\
        fprintf(stderr,"CUDA call failed!\n%s\n",cudaGetErrorString(err)); \
	assert(false);							\
      }									\
  }

#define INDEX(i,j,k)				\
  ((i)*MAX_Y*MAX_Z + (j)*MAX_Z + (k))


inline void swap(int &a, int &b){
  int temp = a;
  a = b; 
  b = temp;
}

//////////////////////////
//FUNCTION DECLARATION
__global__ void initializePressure(int sel);
__global__ void initializeGrid(int sel);
__global__ void initializeGrid(float4* data);

__global__ void force();
__global__ void jacobi(int dst, int src, int field, float a, float invb);
__global__ void jacobiStuff(float4* next, float4* old, float4* src, float4 a,float4 invb);
__global__ void jacobiPressure(int dst, int src);

__global__ void advect(int dst, int src, int vel, float dt);
__global__ void advectStuff(float4* dst, float4* src, int vel, float dt);

__global__ void diverge(int src);
__global__ void project(int dst, int src, int press);

__global__ void sampleField(float4* dst, float4* dst_stuff, int src, int press, 
			    float dt, int inject_status, float4 inject_color);

__global__ void scenario1(int x, int y, int z, 
			  float vx, float vy, float vz, 
			  int src);
__global__ void scenario2(int x0, int y0, int z0, 
			  int x1, int y1, int z1, 
			  float vx0, float vy0, float vz0, 
			  float vx1, float vy1, float vz1, 
			  int src);



__global__ void pressureGradCopy(float4* dst, int src);
__global__ void pressureCopy(float4* dst,  int src);
__global__ void velCopy(float4* dst, int src);
__global__ void divergenceCopy(float4* dst);
/////////////////////////
//DATA DECLARATION
enum BufferID{FIRST=0, SECOND, SWAP, TOTAL};
__device__ float divergence[MAX_X][MAX_Y][MAX_Z];
__device__ float vel_x[TOTAL][MAX_X][MAX_Y][MAX_Z];
__device__ float vel_y[TOTAL][MAX_X][MAX_Y][MAX_Z];
__device__ float vel_z[TOTAL][MAX_X][MAX_Y][MAX_Z];
__device__ float pressure[TOTAL][MAX_X][MAX_Y][MAX_Z];

__device__ float4 d_h_pos;
__device__ float3 d_h_vel;

float4* d_force_array;
float4* h_force_array;

//Rendering color?
float4* stuff[TOTAL];
float4* display_stuff;

//Rendering coordinates
float3 grid[MAX_X][MAX_Y][MAX_Z];

int _src = FIRST;
int _dst = SECOND;
int _swap = SWAP;

int _stuff_src = FIRST;
int _stuff_dst = SECOND;
int _stuff_swap = SWAP;

int _pressure_src = FIRST;
int _pressure_dst = SECOND;

////////////////////////
//RENDER DECLARATION
GLuint fluid_vertexShader;
GLuint fluid_pixelShader;
GLuint fluid_glProgram;
GLuint grid_bo;
GLuint display_bo;
GLuint fluid_texture;

cudaGraphicsResource* display_resource;

extern float3 h_pos;
extern float3 h_force;
extern float3 h_vel;
extern int buttonStatus;
extern bool resetVel;
extern bool resetStuff;
extern float adaptiveDT;


extern float tool_size;
extern float h_color_r;
extern float h_color_g;
extern float h_color_b;
extern float injection_rate;
extern float point_size;
extern float diff_rate;
extern float stuff_diff_rate;

extern int preset;



/////////////////////////
__host__ unsigned char* createGaussianMap(int N);

__host__ void initCUDA(){
  //CUDA initializaiton
  CUDA_SAFE_CALL(cudaGLSetGLDevice(0));

  //the color render buffer
  glGenBuffers(1, &display_bo);
  glBindBuffer(GL_ARRAY_BUFFER_ARB, display_bo);
  glBufferData(GL_ARRAY_BUFFER_ARB, sizeof(float4)*MAX_X*MAX_Y*MAX_Z, 0, GL_DYNAMIC_DRAW_ARB);
  glBindBuffer(GL_ARRAY_BUFFER_ARB, 0);
  cutilSafeCall(cudaGraphicsGLRegisterBuffer(&display_resource, 
					     display_bo,
					     cudaGraphicsMapFlagsWriteDiscard));


  for(int i = 0; i<TOTAL; i++){
    cutilSafeCall(cudaMalloc((void**)&(stuff[i]),  sizeof(float4)*MAX_X*MAX_Y*MAX_Z));
    initializeGrid<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(stuff[i]);
    initializeGrid<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(i);
    initializePressure<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(i);
  }

  //Haptic bufffer
  cutilSafeCall(cudaMalloc((void**)&d_force_array, sizeof(float4)*int(tool_size*tool_size*tool_size*8)));
  h_force_array = (float4*)malloc(sizeof(float4)*int(tool_size*tool_size*tool_size*8));
  printf("CUDA Started\n");
}


__host__ void initGraphics(){
  adaptiveDT = 0.005;

  //OPENGL Vertex buffer
  for(int i = 0; i<MAX_X; i++){
    for(int j = 0; j<MAX_Y; j++){
      for(int k = 0; k<MAX_Z; k++){
	grid[i][j][k] = make_float3(float(i-MAX_X/2), float(j-MAX_Y/2), float(k-MAX_Z/2));
      }
    }
  }
  glGenBuffers(1, &grid_bo);
  glBindBuffer(GL_ARRAY_BUFFER_ARB, grid_bo);
  glBufferData(GL_ARRAY_BUFFER_ARB, sizeof(float3)*MAX_X*MAX_Y*MAX_Z, grid, GL_STATIC_DRAW);
  int size = 0;
  glGetBufferParameteriv(GL_ARRAY_BUFFER_ARB, GL_BUFFER_SIZE, (GLint*)&size); 
  glBindBuffer(GL_ARRAY_BUFFER_ARB, 0);


  //test rendering texture
  int resolution = 32;
  unsigned char* data = createGaussianMap(resolution);
  glGenTextures(1, &fluid_texture);
  glBindTexture(GL_TEXTURE_2D, fluid_texture);
  glTexParameteri(GL_TEXTURE_2D, GL_GENERATE_MIPMAP_SGIS, GL_TRUE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, resolution, resolution, 0, 
	       GL_RGBA, GL_UNSIGNED_BYTE, data);

  //GLSL programs
  fluid_vertexShader = glCreateShader(GL_VERTEX_SHADER);
  fluid_pixelShader = glCreateShader(GL_FRAGMENT_SHADER);
  fluid_glProgram = glCreateProgram();
  string v;
  readSource(v,"point.vert.glsl");
  string p;
  readSource(p,"point.frag.glsl");
  const char* v_c = v.c_str();
  const char* p_c = p.c_str();
  glShaderSource(fluid_vertexShader, 1, &v_c, 0);
  glShaderSource(fluid_pixelShader, 1, &p_c, 0);
  glAttachShader(fluid_glProgram, fluid_vertexShader);
  glAttachShader(fluid_glProgram, fluid_pixelShader);
  glLinkProgram(fluid_glProgram);
  
  printf("OpenGL Started\n");

}

/////////////////////////////
//RENDERING
/////////////////////////////
__host__ void renderCopy(display_mode m){
  size_t mapped_size;
  //register the display buffer
  cutilSafeCall(cudaGraphicsMapResources(1, &display_resource, 0));
  cutilSafeCall(cudaGraphicsResourceGetMappedPointer((void**)&(display_stuff), 
						     &mapped_size, 
						     display_resource));

  switch (m){
  case DISPLAY_FLUID:
    cutilSafeCall(cudaMemcpy(display_stuff, stuff[_stuff_src], mapped_size, cudaMemcpyDeviceToDevice));
    break;
  case DISPLAY_VELOCITY:
    velCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff,_src);
    break;
  case DISPLAY_PRESSURE:
    pressureCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff,_pressure_src);
    break;
  case DISPLAY_PRESSURE_GRAD:
    pressureGradCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff,_pressure_src);
    break;
  case DISPLAY_DIVERGENCE:
    divergenceCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff);
    break;
  default:
    break;
  }
  

  //debugCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff, _pressure_src);
  //debugCopy<<<dim3(MAX_X, MAX_Y), MAX_Z>>>(display_stuff, _pressure_src);

  cutilSafeCall(cudaGraphicsUnmapResources(1, &display_resource, 0));
}

__host__ void renderFluid(){

  // setup point sprites
  glEnable(GL_POINT_SPRITE_ARB);
  glTexEnvi(GL_POINT_SPRITE_ARB, GL_COORD_REPLACE_ARB, GL_TRUE);
  glEnable(GL_VERTEX_PROGRAM_POINT_SIZE_NV);
  glPointSize(point_size);
  glBlendFunc(GL_SRC_ALPHA, GL_ONE);
  glEnable(GL_BLEND);
  glDepthMask(GL_FALSE);
  
  glUseProgram(fluid_glProgram);
  GLuint texLoc = glGetUniformLocation(fluid_glProgram, "splatTexture");
  glUniform1i(texLoc, 0);
  glActiveTextureARB(GL_TEXTURE0_ARB);
  glBindTexture(GL_TEXTURE_2D, fluid_texture);

  glColor3f(1, 1, 1);  
  glSecondaryColor3f(1.f, 0.6f, 0.3f);


  GLint color_att = glGetAttribLocation(fluid_glProgram, "colorIn");
  glBindBuffer(GL_ARRAY_BUFFER, display_bo);
  glEnableVertexAttribArray(color_att);
  glVertexAttribPointer(color_att, 4, GL_FLOAT, false ,0, 0);

  glEnableClientState(GL_VERTEX_ARRAY);                
  glBindBufferARB(GL_ARRAY_BUFFER_ARB, grid_bo);
  glVertexPointer(3, GL_FLOAT, 0, 0);

  glDrawArrays(GL_POINTS, 0, MAX_X*MAX_Y*MAX_Z);
  glBindBufferARB(GL_ARRAY_BUFFER_ARB, 0);
  glDisableClientState(GL_VERTEX_ARRAY); 
  glDisableClientState(GL_COLOR_ARRAY); 
  
  glUseProgram(0);  
  glDisable(GL_POINT_SPRITE_ARB);
  glDisable(GL_BLEND);
  glDepthMask(GL_TRUE);
}
inline float evalHermite(float pA, float pB, float vA, float vB, float u)
{
  float u2=(u*u), u3=u2*u;
  float B0 = 2*u3 - 3*u2 + 1;
  float B1 = -2*u3 + 3*u2;
  float B2 = u3 - 2*u2 + u;
  float B3 = u3 - u;
  return( B0*pA + B1*pB + B2*vA + B3*vB );
}
__host__ unsigned char* createGaussianMap(int N)
{
  float *M = new float[2*N*N];
  unsigned char *B = new unsigned char[4*N*N];
  float X,Y,Y2,Dist;
  float Incr = 2.0f/N;
  int i=0;  
  int j = 0;
  Y = -1.0f;
  //float mmax = 0;
  for (int y=0; y<N; y++, Y+=Incr)
    {
      Y2=Y*Y;
      X = -1.0f;
      for (int x=0; x<N; x++, X+=Incr, i+=2, j+=4)
        {
	  Dist = (float)sqrtf(X*X+Y2);
	  if (Dist>1) Dist=1;
	  M[i+1] = M[i] = evalHermite(1.0f,0,0,0,Dist);
	  B[j+3] = B[j+2] = B[j+1] = B[j] = (unsigned char)(M[i] * 255);
        }
    }
  delete [] M;
  return(B);
}    



///////////////////////////////
//Host compute
//////////////////////////////
__host__ void stepFluid(){
  
  float4 real_pos = make_float4(h_pos) + make_float4(float(MAX_X>>1)-tool_size+1, 
					float(MAX_Y>>1)-tool_size+1,
					float(MAX_Z>>1)-tool_size+1, 
					tool_size);
  cutilSafeCall( cudaMemcpyToSymbol("d_h_pos", &real_pos, sizeof(float4),0,cudaMemcpyHostToDevice ));
  cutilSafeCall( cudaMemcpyToSymbol("d_h_vel", &h_vel, sizeof(float3),0,cudaMemcpyHostToDevice  ));
  
  if(resetStuff){
    resetStuff = false;
    for(int i = 0; i<TOTAL; i++)
      initializeGrid<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(stuff[i]);
  }
  if(resetVel){
    resetVel = false;
    for(int i = 0; i<TOTAL; i++){
      initializeGrid<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(i);
      initializePressure<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(i);
    }
  }

  //Velocity field update
  if(vel_diffusion_enable){
    float alpha = adaptiveDT*diff_rate;

    jacobi<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_dst, _src, _src, alpha, 1/(1+6.0f*alpha));
      
    for(unsigned int i = 0; i<JITERATION; i++){
      swap(_swap, _dst);
      jacobi<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_dst, _swap, _src, alpha, 1/(1+6.0f*alpha));
    }
    swap(_dst, _src);
  }
  if(vel_advect_enable){
    advect<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_dst, _src, _src, adaptiveDT);  
    swap(_dst, _src);
  }
  diverge<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_src);
  initializePressure<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_pressure_src);
  if(vel_projection_enable){
    for(unsigned int i = 0; i<JITERATION*4; i++){
      jacobiPressure<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_pressure_dst, _pressure_src);  
      swap(_pressure_dst, _pressure_src);
    }
    project<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(_dst, _src, _pressure_src);
    swap(_dst, _src);
  }




  //fluid particle movement
  if(density_diffusion_enable){
    float stuff_alpha = stuff_diff_rate*adaptiveDT;
      jacobiStuff<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(stuff[_stuff_dst], stuff[_stuff_src], stuff[_stuff_src], make_float4(stuff_alpha), make_float4(1.f/(1.f+6.f*stuff_alpha)));
    for(unsigned int i = 0; i<JITERATION; i++){
      swap(_stuff_swap, _stuff_dst);
      jacobiStuff<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(stuff[_stuff_dst], stuff[_stuff_swap], stuff[_stuff_src], make_float4(stuff_alpha), make_float4(1.f/(1.f+6.f*stuff_alpha)));
    }
    swap(_stuff_dst, _stuff_src);
  }
  if(density_advect_enable){
    advectStuff<<<dim3(MAX_X,MAX_Y),MAX_Z>>>(stuff[_stuff_dst], stuff[_stuff_src], _src , adaptiveDT);  
    swap(_stuff_dst, _stuff_src);
  }  


  

  
  //General haptic interaction
  if(haptic_effect_enable){
    int edge_length = int(tool_size*2);
    sampleField<<<dim3(edge_length ,edge_length),edge_length >>>(d_force_array, 
								 stuff[_stuff_src], //direclty edit the source
								 _src, //direclty edit the source
								 _pressure_src,//direclty edit the source
								 adaptiveDT,
								 buttonStatus, 
								 TRACE_COLOR);
    cutilSafeCall(cudaMemcpy((void*)h_force_array, 
			     (const void*)d_force_array, 
			     sizeof(float4)*int(edge_length *edge_length *edge_length),
			     cudaMemcpyDeviceToHost));
    h_force.x=0.f;
    h_force.y=0.f;
    h_force.z=0.f;
    for(int i = 0; i<edge_length*edge_length*edge_length; i++){
      h_force.x+=h_force_array[i].x;
      h_force.y+=h_force_array[i].y;
      h_force.z+=h_force_array[i].z;
    }
    h_force.x/=(edge_length *edge_length *edge_length*1.2);
    h_force.y/=(edge_length *edge_length *edge_length*1.2);
    h_force.z/=(edge_length *edge_length *edge_length*1.2);
  }


  if(preset == 1){
    scenario1<<<dim3(8,2), 8 >>>(MAX_X/2 , 2, MAX_Z/2,
				0.f, 30.f, 0.f,
				_src);
  } else if(preset == 2){
    scenario2<<<dim3(4,2), 4 >>>(MAX_X/2 , 2, MAX_Z/2,
				 2, MAX_Y/2,  MAX_Z/2,
				 0.f, 30.f, 0.f,
				 30.f, 0.f, 0.f,
				 _src);
  }else if(preset == 3){
    scenario1<<<dim3(2,4), 2 >>>(MAX_X/2 , 2, MAX_Z/2,
				 0.f, 60.f, 0.f,
				 _src);
  } else if(preset == 4){
    scenario2<<<dim3(4,4), 4 >>>(MAX_X/2-2 , 2, MAX_Z/2-2,
				 2, MAX_Y/2-2,  MAX_Z/2-2,
				 20.f, 20.f, 0.f,
				 20.f,-20.f, 0.f,
				 _src);
    scenario2<<<dim3(4,4), 4 >>>(MAX_X/2-2, MAX_Y-5, MAX_Z/2-2,
				 MAX_X-5 , MAX_Y/2-2, MAX_Z/2-2,
				 -20.f,-20.f, 0.f,
				 -20.f, 20.f, 0.f,
				 _src);
    
  }

  
}







///////////////////////////
//Devie Compute
///////////////////////////

__global__ void initializeGrid(float4* data){
  unsigned int index = INDEX(blockIdx.x, blockIdx.y,threadIdx.x);
  data[index] = make_float4(0.f, 0.f, 0.0f, 1.f);
}


__global__ void initializePressure(int sel){
  pressure[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.0f;    
}
__global__ void initializeGrid(int sel){
  /*
  vel_x[sel][blockIdx.x][blockIdx.y][threadIdx.x] =-INIT_VEL(blockIdx.y);
  vel_y[sel][blockIdx.x][blockIdx.y][threadIdx.x] = INIT_VEL(blockIdx.x);
  vel_z[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.0f;
  pressure[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.0f;

  */

  vel_x[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.f;
  vel_y[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.f;
  vel_z[sel][blockIdx.x][blockIdx.y][threadIdx.x] = 0.0f;
  
}


//regular access
__global__ void force(){

}




//regular access
__global__ void jacobi(int next, int old, int src, float a, float invb){
  /*
  bool pred_x_up  =(blockIdx.x ==MAX_X-1) || inObject(blockIdx.x+1, blockIdx.y, threadIdx.x);
  bool pred_x_down=(blockIdx.x ==0) || inObject(blockIdx.x-1, blockIdx.y, threadIdx.x);
  bool pred_y_up  =(blockIdx.y ==MAX_Y-1) || inObject(blockIdx.x, blockIdx.y+1, threadIdx.x);
  bool pred_y_down=(blockIdx.y ==0) || inObject(blockIdx.x, blockIdx.y-1, threadIdx.x);
  bool pred_z_up  =(threadIdx.x==MAX_Z-1) || inObject(blockIdx.x, blockIdx.y, threadIdx.x+1);
  bool pred_z_down=(threadIdx.x==0) || inObject(blockIdx.x+1, blockIdx.y, threadIdx.x-1);
  */
  bool pred_x_up  =(blockIdx.x ==MAX_X-1);
  bool pred_x_down=(blockIdx.x ==0);
  bool pred_y_up  =(blockIdx.y ==MAX_Y-1);
  bool pred_y_down=(blockIdx.y ==0);
  bool pred_z_up  =(threadIdx.x==MAX_Z-1);
  bool pred_z_down=(threadIdx.x==0);

  int index_x_up  =blockIdx.x +1;//(pred_x_up  ?blockIdx.x :blockIdx.x +1 );
  int index_x_down=blockIdx.x -1;//(pred_x_down?blockIdx.x :blockIdx.x -1 );
  int index_y_up  =blockIdx.y +1;//(pred_y_up  ?blockIdx.y :blockIdx.y +1 );
  int index_y_down=blockIdx.y -1;//(pred_y_down?blockIdx.y :blockIdx.y -1 );
  int index_z_up  =threadIdx.x+1;//(pred_z_up  ?threadIdx.x:threadIdx.x+1 );
  int index_z_down=threadIdx.x-1;//(pred_z_down?threadIdx.x:threadIdx.x-1 );

  float x_bound = vel_x[old][blockIdx.x][blockIdx.y][threadIdx.x];
  float x_z_up  = (pred_z_up  ? x_bound:vel_x[old][blockIdx.x  ][blockIdx.y  ][index_z_up   ]); 
  float x_z_down= (pred_z_down? x_bound:vel_x[old][blockIdx.x  ][blockIdx.y  ][index_z_down ]);
  float x_y_up  = (pred_y_up  ? x_bound:vel_x[old][blockIdx.x  ][index_y_up  ][threadIdx.x  ]); 
  float x_y_down= (pred_y_down? x_bound:vel_x[old][blockIdx.x  ][index_y_down][threadIdx.x  ]);
  float x_x_up  = (pred_x_up  ?-x_bound:vel_x[old][index_x_up  ][blockIdx.y  ][threadIdx.x  ]); 
  float x_x_down= (pred_x_down?-x_bound:vel_x[old][index_x_down][blockIdx.y  ][threadIdx.x  ]);

  float y_bound = vel_y[old][blockIdx.x][blockIdx.y][threadIdx.x];
  float y_z_up  = (pred_z_up  ? y_bound:vel_y[old][blockIdx.x  ][blockIdx.y  ][index_z_up   ]); 
  float y_z_down= (pred_z_down? y_bound:vel_y[old][blockIdx.x  ][blockIdx.y  ][index_z_down ]);
  float y_y_up  = (pred_y_up  ?-y_bound:vel_y[old][blockIdx.x  ][index_y_up  ][threadIdx.x  ]); 
  float y_y_down= (pred_y_down?-y_bound:vel_y[old][blockIdx.x  ][index_y_down][threadIdx.x  ]);
  float y_x_up  = (pred_x_up  ? y_bound:vel_y[old][index_x_up  ][blockIdx.y  ][threadIdx.x  ]); 
  float y_x_down= (pred_x_down? y_bound:vel_y[old][index_x_down][blockIdx.y  ][threadIdx.x  ]);

  float z_bound = vel_z[old][blockIdx.x][blockIdx.y][threadIdx.x];
  float z_z_up  = (pred_z_up  ?-z_bound:vel_z[old][blockIdx.x  ][blockIdx.y  ][index_z_up   ]); 
  float z_z_down= (pred_z_down?-z_bound:vel_z[old][blockIdx.x  ][blockIdx.y  ][index_z_down ]);
  float z_y_up  = (pred_y_up  ? z_bound:vel_z[old][blockIdx.x  ][index_y_up  ][threadIdx.x  ]); 
  float z_y_down= (pred_y_down? z_bound:vel_z[old][blockIdx.x  ][index_y_down][threadIdx.x  ]);
  float z_x_up  = (pred_x_up  ? z_bound:vel_z[old][index_x_up  ][blockIdx.y  ][threadIdx.x  ]); 
  float z_x_down= (pred_x_down? z_bound:vel_z[old][index_x_down][blockIdx.y  ][threadIdx.x  ]);


  vel_x[next][blockIdx.x][blockIdx.y][threadIdx.x] =((x_z_up+x_z_down+
						      x_y_up+x_y_down+
						      x_x_up+x_x_down)*a+
						     vel_x[src][blockIdx.x][blockIdx.y][threadIdx.x])*invb;
  vel_y[next][blockIdx.x][blockIdx.y][threadIdx.x] =((y_z_up+y_z_down+
						      y_y_up+y_y_down+
						      y_x_up+y_x_down)*a+
						     vel_y[src][blockIdx.x][blockIdx.y][threadIdx.x])*invb;
  vel_z[next][blockIdx.x][blockIdx.y][threadIdx.x] =((z_z_up+z_z_down+
						      z_y_up+z_y_down+
						      z_x_up+z_x_down)*a+
						     vel_z[src][blockIdx.x][blockIdx.y][threadIdx.x])*invb;
  
}


__global__ void jacobiStuff(float4* next, float4* old, float4* src, float4 a,float4 invb){
  unsigned int index = INDEX(blockIdx.x, blockIdx.y,  threadIdx.x);
  
  int x_up_index  =min(int(blockIdx.x) +1,MAX_X-1);
  int x_down_index=max(int(blockIdx.x) -1,0);
  int y_up_index  =min(int(blockIdx.y) +1,MAX_Y-1);
  int y_down_index=max(int(blockIdx.y) -1,0);
  int z_up_index  =min(int(threadIdx.x)+1,MAX_Z-1);
  int z_down_index=max(int(threadIdx.x)-1,0);
  
  //float4 bound = old[index];
  
  float4 z_up  = old[INDEX(blockIdx.x,blockIdx.y,z_up_index  )]; 
  float4 z_down= old[INDEX(blockIdx.x,blockIdx.y,z_down_index)];
  float4 y_up  = old[INDEX(blockIdx.x,y_up_index  ,threadIdx.x)]; 
  float4 y_down= old[INDEX(blockIdx.x,y_down_index,threadIdx.x)];
  float4 x_up  = old[INDEX(x_up_index  ,blockIdx.y,threadIdx.x)]; 
  float4 x_down= old[INDEX(x_down_index,blockIdx.y,threadIdx.x)];
 
  next[index] =((z_up+z_down+
		 y_up+y_down+
		 x_up+x_down)*a+src[index])*invb;
}

//laplacian of the pressure = diveragence of velocity
__global__ void jacobiPressure(int dst, int src){
  int x_up_index  =min(int(blockIdx.x) +1,MAX_X-1);
  int x_down_index=max(int(blockIdx.x) -1,0);
  int y_up_index  =min(int(blockIdx.y) +1,MAX_Y-1);
  int y_down_index=max(int(blockIdx.y) -1,0);
  int z_up_index  =min(int(threadIdx.x)+1,MAX_Z-1);
  int z_down_index=max(int(threadIdx.x)-1,0);

  float x_up = pressure[src][x_up_index  ][blockIdx.y  ][threadIdx.x ];
  float x_down=pressure[src][x_down_index][blockIdx.y  ][threadIdx.x ];
  float y_up = pressure[src][blockIdx.x  ][y_up_index  ][threadIdx.x ];
  float y_down=pressure[src][blockIdx.x  ][y_down_index][threadIdx.x ];
  float z_up = pressure[src][blockIdx.x  ][blockIdx.y  ][z_up_index  ];
  float z_down=pressure[src][blockIdx.x  ][blockIdx.y  ][z_down_index]; 

  pressure[dst][blockIdx.x][blockIdx.y][threadIdx.x] =
    (x_up + x_down +
     y_up + y_down +
     z_up + z_down -
     divergence[blockIdx.x][blockIdx.y][threadIdx.x])*1.f/6.0f;
}


__global__ void advect(int dst, int src, int vel, float dt){
  //reverse time 
  float x = clamp((float(blockIdx.x) -dt*vel_x[vel][blockIdx.x][blockIdx.y][threadIdx.x]),-0.5f,float(MAX_X-1)+0.5f);
  float y = clamp((float(blockIdx.y) -dt*vel_y[vel][blockIdx.x][blockIdx.y][threadIdx.x]),-0.5f,float(MAX_Y-1)+0.5f);
  float z = clamp((float(threadIdx.x)-dt*vel_z[vel][blockIdx.x][blockIdx.y][threadIdx.x]),-0.5f,float(MAX_Z-1)+0.5f);
  //trilinear interpolation
  int i = (int)x;
  int j = (int)y;
  int k = (int)z;
  int i_p = min(int(x+1.f), MAX_X-1);
  int j_p = min(int(y+1.f), MAX_Y-1);
  int k_p = min(int(z+1.f), MAX_Z-1);

  float a= x-i;
  float b= y-j;
  float c= z-k;

  float x_scale = ((x<0.f) || (x>float(MAX_X-1)))?0.f:1.f;
  float y_scale = ((y<0.f) || (y>float(MAX_Y-1)))?0.f:1.f;
  float z_scale = ((z<0.f) || (z>float(MAX_Z-1)))?0.f:1.f;

  vel_x[dst][blockIdx.x][blockIdx.y][threadIdx.x] = 
    lerp(
	 lerp(
	      lerp(vel_x[src][i][j]  [k],vel_x[src][i][j]  [k_p], c), 
	      lerp(vel_x[src][i][j_p][k],vel_x[src][i][j_p][k_p] ,c),b), 
	 lerp(
	      lerp(vel_x[src][i_p][j]  [k],vel_x[src][i_p][j]  [k_p],c), 
	      lerp(vel_x[src][i_p][j_p][k],vel_x[src][i_p][j_p][k_p],c),b)  ,a)*x_scale;
  
  vel_y[dst][blockIdx.x][blockIdx.y][threadIdx.x] = 
    lerp(
	 lerp(
	      lerp(vel_y[src][i][j]  [k],vel_y[src][i][j]  [k_p], c), 
	      lerp(vel_y[src][i][j_p][k],vel_y[src][i][j_p][k_p] ,c),b), 
	 lerp(
	      lerp(vel_y[src][i_p][j]  [k],vel_y[src][i_p][j]  [k_p],c), 
	      lerp(vel_y[src][i_p][j_p][k],vel_y[src][i_p][j_p][k_p],c),b)  ,a)*y_scale;
    vel_z[dst][blockIdx.x][blockIdx.y][threadIdx.x] = 
    lerp(
	 lerp(
	      lerp(vel_z[src][i][j]  [k],vel_z[src][i][j]  [k_p], c), 
	      lerp(vel_z[src][i][j_p][k],vel_z[src][i][j_p][k_p] ,c),b), 
	 lerp(
	      lerp(vel_z[src][i_p][j]  [k],vel_z[src][i_p][j]  [k_p],c), 
	      lerp(vel_z[src][i_p][j_p][k],vel_z[src][i_p][j_p][k_p],c),b)  ,a)*z_scale;
  
  
}



__global__ void advectStuff(float4* dst, float4* src, int vel, float dt){

  unsigned int index = INDEX(blockIdx.x, blockIdx.y,  threadIdx.x);
  //reverse time 
  float x = clamp((float(blockIdx.x)  - dt*vel_x[vel][blockIdx.x][blockIdx.y][threadIdx.x]), -0.5f, float(MAX_X-1)+.5f);
  float y = clamp((float(blockIdx.y)  - dt*vel_y[vel][blockIdx.x][blockIdx.y][threadIdx.x]), -0.5f, float(MAX_Y-1)+.5f);
  float z = clamp((float(threadIdx.x) - dt*vel_z[vel][blockIdx.x][blockIdx.y][threadIdx.x]), -0.5f, float(MAX_Z-1)+.5f);

  //trilinear interpolation
  int i = (int)x;
  int j = (int)y;
  int k = (int)z;
  int i_p = min(int(x+1.f), MAX_X-1);
  int j_p = min(int(y+1.f), MAX_Y-1);
  int k_p = min(int(z+1.f), MAX_Z-1);

  float a= abs(x-i);
  float b= abs(y-j);
  float c= abs(z-k);
  
  dst[index] =
    lerp(
	 lerp(
	      lerp(src[INDEX(i,j,k)],src[INDEX(i,j,k_p)], c), 
	      lerp(src[INDEX(i,j_p,k)],src[INDEX(i,j_p,k_p)] ,c),b), 
	 lerp(
	      lerp(src[INDEX(i_p,j,k)],src[INDEX(i_p,j,k_p)],c), 
	      lerp(src[INDEX(i_p,j_p,k)],src[INDEX(i_p,j_p,k_p)],c),b)  ,a);

}



//divergence of the velocity field
__global__ void diverge(int src){
  float x_bound = vel_x[src][blockIdx.x][blockIdx.y][threadIdx.x];
  float y_bound = vel_y[src][blockIdx.x][blockIdx.y][threadIdx.x];
  float z_bound = vel_z[src][blockIdx.x][blockIdx.y][threadIdx.x];


  float x_up = (blockIdx.x==MAX_X-1)?-x_bound:vel_x[src][blockIdx.x+1][blockIdx.y  ][threadIdx.x  ];
  float x_down=(blockIdx.x==0      )?-x_bound:vel_x[src][blockIdx.x-1][blockIdx.y  ][threadIdx.x  ];
  float y_up = (blockIdx.y==MAX_Y-1)?-y_bound:vel_y[src][blockIdx.x  ][blockIdx.y+1][threadIdx.x  ];
  float y_down=(blockIdx.y==0      )?-y_bound:vel_y[src][blockIdx.x  ][blockIdx.y-1][threadIdx.x  ];
  float z_up =(threadIdx.x==MAX_Z-1)?-z_bound:vel_z[src][blockIdx.x  ][blockIdx.y  ][threadIdx.x+1];
  float z_down=(threadIdx.x==0     )?-z_bound:vel_z[src][blockIdx.x  ][blockIdx.y  ][threadIdx.x-1]; 
  divergence[blockIdx.x][blockIdx.y][threadIdx.x] = 
    0.5f* ((x_up - x_down)+
	   (y_up - y_down)+
	   (z_up - z_down));
}

//velocity - gradP
__global__ void project(int dst, int src, int press){

  int x_up_index  =min(int(blockIdx.x) +1,MAX_X-1);
  int x_down_index=max(int(blockIdx.x) -1,0);
  int y_up_index  =min(int(blockIdx.y) +1,MAX_Y-1);
  int y_down_index=max(int(blockIdx.y) -1,0);
  int z_up_index  =min(int(threadIdx.x)+1,MAX_Z-1);
  int z_down_index=max(int(threadIdx.x)-1,0);

  float x_up = pressure[press][x_up_index  ][blockIdx.y  ][threadIdx.x ];
  float x_down=pressure[press][x_down_index][blockIdx.y  ][threadIdx.x ];
  float y_up = pressure[press][blockIdx.x  ][y_up_index  ][threadIdx.x ];
  float y_down=pressure[press][blockIdx.x  ][y_down_index][threadIdx.x ];
  float z_up = pressure[press][blockIdx.x  ][blockIdx.y  ][z_up_index  ];
  float z_down=pressure[press][blockIdx.x  ][blockIdx.y  ][z_down_index]; 

  float pgrad_x=(x_up-x_down)*0.5f;
  float pgrad_y=(y_up-y_down)*0.5f;
  float pgrad_z=(z_up-z_down)*0.5f;

  vel_x[dst][blockIdx.x][blockIdx.y][threadIdx.x] = vel_x[src][blockIdx.x][blockIdx.y][threadIdx.x] - pgrad_x;
  vel_y[dst][blockIdx.x][blockIdx.y][threadIdx.x] = vel_y[src][blockIdx.x][blockIdx.y][threadIdx.x] - pgrad_y;
  vel_z[dst][blockIdx.x][blockIdx.y][threadIdx.x] = vel_z[src][blockIdx.x][blockIdx.y][threadIdx.x] - pgrad_z;
}


__global__ void sampleField(float4* dst,  //force
			    float4* dst_stuff, //add color
			    int src, //vel
			    int press, //pressure
			    float dt,  int inject_status, float4 inject_color){
  unsigned int index = blockIdx.x*blockDim.x*blockDim.x + blockIdx.y*blockDim.x+threadIdx.x;
  float x = d_h_pos.x+float(blockIdx.x);
  float y = d_h_pos.y+float(blockIdx.y);
  float z = d_h_pos.z+float(threadIdx.x);

  int i = int(x);
  int j = int(y);
  int k = int(z);

  dst[index]=make_float4(0.f);
  if(i<0||i>MAX_X-1 ||
     j<0||j>MAX_Y-1 ||
     k<0||k>MAX_Z-1){
    return;
  } 
  if(blockIdx.x==0 ){
    dst[index].x += (vel_x[src][i][j][k]-d_h_vel.x);
    vel_x[src][i][j][k]=d_h_vel.x;
  }
  if(blockIdx.y==0  ){
    dst[index].y += (vel_y[src][i][j][k]-d_h_vel.y);
    vel_y[src][i][j][k]=d_h_vel.y;
  }
  if(threadIdx.x==0){
    dst[index].z += (vel_z[src][i][j][k]-d_h_vel.z);
    vel_z[src][i][j][k]=d_h_vel.z;
  }

  if(blockIdx.x==blockDim.x-1){
    dst[index].x += (vel_x[src][i][j][k]-d_h_vel.x);
    vel_x[src][i][j][k]=d_h_vel.x;
  }
  if(blockIdx.y==blockDim.x-1){
    dst[index].y += (vel_y[src][i][j][k]-d_h_vel.y);
    vel_y[src][i][j][k]=d_h_vel.y;
  }
  if(threadIdx.x==blockDim.x-1){
    dst[index].z += (vel_z[src][i][j][k]-d_h_vel.z);
    vel_z[src][i][j][k]=d_h_vel.z;
  }


  dst_stuff[INDEX(i,j,k)] +=inject_color*((inject_status&0x1)?dt:0.f);
}



__global__ void scenario1(int x, int y, int z, 
			 float vx, float vy, float vz, 
			 int src){
    int i = x+blockIdx.x;
    int j = y+blockIdx.y;
    int k = z+threadIdx.x;
    vel_x[src][i][j][k] = vx;
    vel_y[src][i][j][k] = vy;
    vel_z[src][i][j][k] = vz;
}
__global__ void scenario2(int x0, int y0, int z0, 
			  int x1, int y1, int z1, 
			  float vx0, float vy0, float vz0, 
			  float vx1, float vy1, float vz1, 
			  int src){
    int i = x0+blockIdx.x;
    int j = y0+blockIdx.y;
    int k = z0+threadIdx.x;
    vel_x[src][i][j][k] = vx0;
    vel_y[src][i][j][k] = vy0;
    vel_z[src][i][j][k] = vz0;

    i = x1+blockIdx.x;
    j = y1+blockIdx.y;
    k = z1+threadIdx.x;
    vel_x[src][i][j][k] = vx1;
    vel_y[src][i][j][k] = vy1;
    vel_z[src][i][j][k] = vz1;
}


///////////////////
//device debug

__global__ void pressureGradCopy(float4* dst, int src){

  unsigned int index = INDEX(blockIdx.x, blockIdx.y, threadIdx.x);
  float4 lol;
  int x_up_index  =min(int(blockIdx.x) +1,MAX_X-1);
  int x_down_index=max(int(blockIdx.x) -1,0);
  int y_up_index  =min(int(blockIdx.y) +1,MAX_Y-1);
  int y_down_index=max(int(blockIdx.y) -1,0);
  int z_up_index  =min(int(threadIdx.x)+1,MAX_Z-1);
  int z_down_index=max(int(threadIdx.x)-1,0);

  float x_up = pressure[src][x_up_index  ][blockIdx.y  ][threadIdx.x ];
  float x_down=pressure[src][x_down_index][blockIdx.y  ][threadIdx.x ];
  float y_up = pressure[src][blockIdx.x  ][y_up_index  ][threadIdx.x ];
  float y_down=pressure[src][blockIdx.x  ][y_down_index][threadIdx.x ];
  float z_up = pressure[src][blockIdx.x  ][blockIdx.y  ][z_up_index  ];
  float z_down=pressure[src][blockIdx.x  ][blockIdx.y  ][z_down_index]; 
  lol.x=abs((x_up-x_down)*0.5f);
  lol.y=abs((y_up-y_down)*0.5f);
  lol.z=abs((z_up-z_down)*0.5f);
  dst[index]=lol;
}

__global__ void pressureCopy(float4* dst,  int src){
  unsigned int index = INDEX(blockIdx.x, blockIdx.y, threadIdx.x);
  dst[index]=make_float4(pressure[src][blockIdx.x][blockIdx.y][threadIdx.x]*10);
}

__global__ void velCopy(float4* dst, int src){
  unsigned int index = INDEX(blockIdx.x, blockIdx.y, threadIdx.x);
  float4 lol;
  lol.x=abs(vel_x[src][blockIdx.x][blockIdx.y][threadIdx.x ]/(MAX_X*4));
  lol.y=abs(vel_y[src][blockIdx.x][blockIdx.y][threadIdx.x ]/(MAX_Y*4));
  lol.z=abs(vel_z[src][blockIdx.x][blockIdx.y][threadIdx.x ]/(MAX_Z*4));
  normalize(lol);
  dst[index]=lol;
}


__global__ void divergenceCopy(float4* dst ){
  unsigned int index = INDEX(blockIdx.x, blockIdx.y, threadIdx.x);
  dst[index]=make_float4(divergence[blockIdx.x][blockIdx.y][threadIdx.x]);
}

